# 整合版：使用Neo4j构建Text2SQL语义图谱
from camel.agents import ChatAgent
from camel.messages import BaseMessage
from camel.models import ModelFactory
from camel.types import ModelPlatformType, RoleType
# 1. 导入依赖库
from py2neo import Graph
from dotenv import load_dotenv
import pymysql

load_dotenv()
import os

# 2. 数据库配置（请根据实际环境修改）
MYSQL_CONFIG = {
    "host": "localhost",
    "port": 3306,  # MySQL默认端口，根据实际情况修改
    "user": os.getenv('MYSQL_USER'),
    "password": os.getenv('MYSQL_PASSWORD'),  # 替换为你的MySQL密码
    "database": os.getenv('DATABASE'),  # 数据库名称
    "charset": "utf8mb4",
}

NEO4J_CONFIG = {
    "uri": "bolt://localhost:7687",  # Neo4j默认Bolt端口
    "user": "neo4j",  # Neo4j默认用户名
    "password": os.getenv('PASSWORD_NEO'),  # 替换为你的Neo4j密码
}

# 3. 表间关系定义（人工补充业务语义）
TABLE_RELATIONSHIPS = [
    {
        "from_table": "t_customers",
        "to_table": "t_sales_orders",
        "description": "客户创建销售订单",
        "field_relation": "customer_id → order.customer_id",
    },
    {
        "from_table": "t_sales_orders",
        "to_table": "t_order_details",
        "description": "订单包含多个明细项",
        "field_relation": "order_id → detail.order_id",
    },
    {
        "from_table": "t_products",
        "to_table": "t_order_details",
        "description": "产品属于订单明细",
        "field_relation": "product_id → detail.product_id",
    },
    {
        "from_table": "t_user",
        "to_table": "t_user_qa_record",
        "description": "用户产生问答记录",
        "field_relation": "id → record.user_id",
    },
]


# 4. MySQL操作函数
def connect_mysql():
    """连接MySQL数据库并返回连接对象"""
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        print("✅ 成功连接到MySQL数据库")
        return conn
    except Exception as e:
        raise Exception(f"MySQL连接失败: {str(e)}")

# 5. Neo4j操作函数
def connect_neo4j():
    """连接Neo4j图数据库并返回Graph对象"""
    try:
        graph = Graph(
            uri=NEO4J_CONFIG["uri"],
            auth=(NEO4J_CONFIG["user"], NEO4J_CONFIG["password"])
        )
        print("✅ 成功连接到Neo4j数据库")
        return graph
    except Exception as e:
        raise Exception(f"Neo4j连接失败: {str(e)}")

def get_table_structure(mysql_conn):
    """从MySQL数据库提取表结构信息"""
    tables = {}

    with mysql_conn.cursor() as cursor:
        # 获取所有表名
        cursor.execute("SHOW TABLES")
        table_names = [row[0] for row in cursor.fetchall()]

        if not table_names:
            raise Exception("数据库中未找到任何表，请先执行建表脚本")

        # 提取每个表的字段信息
        for table_name in table_names:
            cursor.execute(f"SHOW COLUMNS FROM {table_name}")
            columns = cursor.fetchall()

            fields = []
            for col in columns:
                field_name = col[0]
                key_type = col[3]  # 字段的键类型：PRI=主键，MUL=外键/索引

                # 标记主键和外键，方便后续图谱理解
                if key_type == "PRI":
                    fields.append(f"{field_name} [主键]")
                elif key_type == "MUL":
                    fields.append(f"{field_name} [外键]")
                else:
                    fields.append(field_name)

            tables[table_name] = {
                "name": table_name,
                "fields": fields
            }

    print(f"✅ 成功提取 {len(tables)} 张表的结构信息")
    return tables




def create_neo4j_constraints(graph):
    """创建Neo4j节点唯一性约束，防止表节点重复"""
    # 为Table节点的name属性创建唯一约束
    graph.run("CREATE CONSTRAINT IF NOT EXISTS FOR (t:Table) REQUIRE t.name IS UNIQUE")
    """
    CREATE CONSTRAINT IF NOT EXISTS：创建约束，如果已存在则跳过
    FOR (t:Table)：指定约束作用于标签为Table的节点
    REQUIRE t.name IS UNIQUE：要求节点的name属性值必须唯一
    """
    print("✅ 已创建Neo4j节点唯一性约束")


def create_table_nodes(graph, tables):
    """在Neo4j中创建表节点，存储表名和字段信息
    这个函数将MySQL表结构信息转换为Neo4j节点，每个数据库表对应一个Table节点。
    """
    for table_info in tables.values():
        # 使用MERGE确保节点唯一，避免重复创建
        graph.run(
            """
            MERGE (t:Table {name: $name})
            SET t.label = $label, t.fields = $fields
            """,
            name=table_info["name"],
            label=table_info["name"],
            fields=table_info["fields"]
        )
        """
        MERGE (t:Table {name: $name})：如果不存在指定name的Table节点则创建，存在则匹配
        SET t.label = $label, t.fields = $fields：设置节点的属性值
        """

    print(f"✅ 已在Neo4j中创建 {len(tables)} 个表节点")


def create_relationship_edges(graph):
    """在Neo4j中创建表之间的关系边"""
    for rel in TABLE_RELATIONSHIPS:
        # 匹配源表和目标表节点，创建关系
        graph.run(
            """
            MATCH (from:Table {name: $from_table})
            MATCH (to:Table {name: $to_table})
            MERGE (from)-[r:REFERENCES {
                description: $description,
                field_relation: $field_relation
            }]->(to)
            """,
            from_table=rel["from_table"],
            to_table=rel["to_table"],
            description=rel["description"],
            field_relation=rel["field_relation"]
        )

    print(f"✅ 已在Neo4j中创建 {len(TABLE_RELATIONSHIPS)} 条表关系")


def query_graph_relationships(graph, table_names=None):
    """查询Neo4j中表之间的关系,
    这个函数的作用是从Neo4j图数据库中查询指定表之间的关系信息，为Text2SQL系统提供表间关联的语义信息。
    """
    if not table_names:   #如果未指定要查询的表名，则使用默认表名列表
        table_names = ["t_customers", "t_sales_orders", "t_products", "t_order_details"]

    query = """
    MATCH (t1:Table)-[r:REFERENCES]-(t2:Table)
    WHERE t1.name IN $table_names AND t2.name IN $table_names AND t1.name < t2.name
    RETURN t1.name AS from_table, r.field_relation AS relationship, t2.name AS to_table
    """

    result = graph.run(query, table_names=table_names).data()
    print('✅ 成功获取到Neo4j中表之间的关系')
    return result


# 大模型系统提示（强化格式要求）
SYSTEM_PROMPT = """
你是Text2SQL专家，必须严格按以下规则将用户问题转为MySQL语句：

### 一、工具调用强制格式（必须严格遵守，否则无法执行）
1. 查询所有表结构：
   [query_table_structure()]
2. 查询指定表结构（例如产品表和订单明细表）：
   [query_table_structure(tables=["t_products", "t_order_details"])]
3. 查询所有表关系：
   [query_table_relationships()]
4. 查询指定表关系（例如产品表到订单明细表）：
   [query_table_relationships(from_table="t_products", to_table="t_order_details")]

### 二、工作流程
1. 第一步：调用query_table_structure()获取所有表名和字段（必须执行）；
2. 第二步：调用query_table_relationships()确认产品表和订单表的关联关系（必须执行）；
3. 第三步：基于工具返回的**实际表名和字段**生成SQL，禁止使用默认表名（如products）。

### 三、SQL生成要求
- 表名和字段名必须与工具返回的完全一致（例如工具返回"t_products"，不可简写）；
- 生成SQL后必须附带说明，解释使用的表、字段和计算逻辑。
"""


# 6. 主函数：构建Text2SQL语义图谱
def build_text2sql_graph(user_question):
    """完整构建流程：从数据库提取信息到Neo4j图谱创建"""
    mysql_conn = None

    try:
        # 连接数据库
        mysql_conn = connect_mysql()
        neo4j_graph = connect_neo4j()

        # 提取MySQL表结构
        table_structure = get_table_structure(mysql_conn)

        # 清空Neo4j中的旧数据（注意：生产环境请谨慎使用）
        print("⚠️ 正在清空Neo4j中的旧数据...")
        neo4j_graph.delete_all()

        # 创建Neo4j约束
        create_neo4j_constraints(neo4j_graph)

        # 创建表节点
        create_table_nodes(neo4j_graph, table_structure)

        # 创建表关系
        create_relationship_edges(neo4j_graph)

        # 验证图谱关系
        # print("\n📊 验证图谱关系：")
        # relationships = query_graph_relationships(neo4j_graph)
        # for rel in relationships:
        #     print(f"   {rel['from_table']} 与 {rel['to_table']} 通过 {rel['relationship']} 关联")

        print("\n🎉 Text2SQL语义图谱构建完成！")
        print("   可通过 http://localhost:7474 访问Neo4j浏览器查看图谱")
        print("   在Neo4j浏览器中执行 MATCH (n) RETURN n 可查看所有节点和关系")
        Model_Type = os.getenv("MODEL_TYPE")
        Model_Api = os.getenv("ZHIPU_API_KEY")
        Model_Url = os.getenv("MODEL_URL")
        # 修复工具调用时MySQL连接对象传递问题，确保get_table_structure能正确获取连接
        model = ModelFactory.create(
            model_platform=ModelPlatformType.OPENAI,
            model_type=Model_Type,
            api_key=Model_Api,
            url=Model_Url
        )

        # 定义包装器，确保get_table_structure始终使用已建立的mysql_conn
        def get_table_structure_tool():
            """
            工具描述：获取所有表的结构信息，包括表名、字段名、字段类型、主键信息等。
            返回值：以字典形式返回所有表的结构详情，适用于数据库结构理解与SQL生成。
            """
            return get_table_structure(mysql_conn)

        # 同理，包装query_graph_relationships，确保传递neo4j_graph
        def query_graph_relationships_tool(table_names=None):
            """
            工具描述：查询指定表之间的关系（如外键、业务关联等），如未指定表名则返回所有表关系。
            参数：table_names（可选）- 需要查询关系的表名列表。
            返回值：返回表与表之间的关系详情，便于理解数据表之间的业务联系。
            """
            return query_graph_relationships(neo4j_graph, table_names)

        agent = ChatAgent(
            model=model,
            system_message=SYSTEM_PROMPT,
            output_language="中文",
            tools=[query_graph_relationships_tool, get_table_structure_tool]
        )
        msg = BaseMessage(
            role_name="用户",
            role_type=RoleType.USER,
            content=user_question,
            meta_dict={}
        )
        response = agent.step(msg)
        print(f"大模型响应：{response.msg.content}")



    except Exception as e:
        print(f"\n❌ 构建过程出错: {str(e)}")
        raise  # 抛出异常以便调试

    finally:
        # 确保MySQL连接被关闭
        if mysql_conn:
            mysql_conn.close()
            print("\n🔌 MySQL连接已关闭")


# 7. 程序入口
if __name__ == "__main__":
    question = "查询每个产品的销售总数量和总金额，并按总金额降序排序"
    build_text2sql_graph(question)
